{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "baba5fab",
   "metadata": {},
   "source": [
    "# Moving Horizon Estimation\n",
    "\n",
    "Richard M. Murray, 24 Feb 2023\n",
    "\n",
    "In this notebook we illustrate the implementation of moving horizon estimation (MHE)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36715c5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import matplotlib.pyplot as plt\n",
    "import control as ct\n",
    "\n",
    "import control.optimal as opt\n",
    "import control.flatsys as fs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d72a155b",
   "metadata": {},
   "source": [
    "## System Description\n",
    "\n",
    "We use the PVTOL dynamics from the textbook, which are contained in the `pvtol` module:\n",
    "\n",
    "<table>\n",
    "<tr>\n",
    "    <td width=\"40%\"><img src=\"https://fbswiki.org/wiki/images/7/76/Pvtol.png\" width=240></td>\n",
    "    <td width=20%></td>\n",
    "    <td width=\"40%\">\n",
    "$$\n",
    "\\begin{aligned}\n",
    "  m \\ddot x &= F_1 \\cos\\theta - F_2 \\sin\\theta - c \\dot x, \\\\\n",
    "  m \\ddot y &= F_1 \\sin\\theta + F_2 \\cos\\theta - m g - c \\dot y, \\\\\n",
    "  J \\ddot \\theta &= r F_1.\n",
    "\\end{aligned}\n",
    "$$\n",
    "    </td>\n",
    "</tr>\n",
    "</table>\n",
    "\n",
    "The measured values of the system are the position and orientation,\n",
    "with added noise $n_x$, $n_y$, and $n_\\theta$:\n",
    "\n",
    "$$\n",
    "  \\vec y = \\begin{bmatrix} x \\\\ y \\\\ \\theta \\end{bmatrix} + \n",
    "  \\begin{bmatrix} n_x \\\\ n_y \\\\ n_z \\end{bmatrix}.\n",
    "$$\n",
    "\n",
    "The parameter values for the PVTOL system come from the Caltech ducted fan experiment, described in more detail in [Lecture 4b](cds112-L4b_pvtol-lqr.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08919988",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pvtol = nominal system (no disturbances or noise)\n",
    "# noisy_pvtol = pvtol w/ process disturbances and sensor noise\n",
    "from pvtol import pvtol, pvtol_noisy, plot_results\n",
    "import pvtol as pvt\n",
    "\n",
    "# Find the equiblirum point corresponding to the origin\n",
    "xe, ue = ct.find_eqpt(\n",
    "    pvtol, np.zeros(pvtol.nstates),\n",
    "    np.zeros(pvtol.ninputs), [0, 0, 0, 0, 0, 0],\n",
    "    iu=range(2, pvtol.ninputs), iy=[0, 1])\n",
    "\n",
    "# Initial condition = 2 meters right, 1 meter up\n",
    "x0, u0 = ct.find_eqpt(\n",
    "    pvtol, np.zeros(pvtol.nstates),\n",
    "    np.zeros(pvtol.ninputs), np.array([2, 1, 0, 0, 0, 0]),\n",
    "    iu=range(2, pvtol.ninputs), iy=[0, 1])\n",
    "\n",
    "# Extract the linearization for use in LQR design\n",
    "pvtol_lin = pvtol.linearize(xe, ue)\n",
    "A, B = pvtol_lin.A, pvtol_lin.B\n",
    "\n",
    "print(pvtol, \"\\n\")\n",
    "print(pvtol_noisy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5771ab93",
   "metadata": {},
   "source": [
    "### Control Design\n",
    "\n",
    "We begin by designing an LQR conroller than can be used for trajectory tracking, which is described in more detail in [Lecture 4b](cds112-L4b_pvtol-lqr.ipynb):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2e88938",
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# LQR design w/ physically motivated weighting\n",
    "#\n",
    "# Shoot for 10 cm error in x, 10 cm error in y.  Try to keep the angle\n",
    "# less than 5 degrees in making the adjustments.  Penalize side forces\n",
    "# due to loss in efficiency.\n",
    "#\n",
    "\n",
    "Qx = np.diag([100, 10, (180/np.pi) / 5, 0, 0, 0])\n",
    "Qu = np.diag([10, 1])\n",
    "K, _, _ = ct.lqr(A, B, Qx, Qu)\n",
    "\n",
    "# Compute the full state feedback solution\n",
    "lqr_ctrl, _ = ct.create_statefbk_iosystem(pvtol, K)\n",
    "\n",
    "# Define the closed loop system that will be used to generate trajectories\n",
    "lqr_clsys = ct.interconnect(\n",
    "    [pvtol_noisy, lqr_ctrl],\n",
    "    inplist = lqr_ctrl.input_labels[0:pvtol.ninputs + pvtol.nstates] + \\\n",
    "        pvtol_noisy.input_labels[pvtol.ninputs:],\n",
    "    inputs = lqr_ctrl.input_labels[0:pvtol.ninputs + pvtol.nstates] + \\\n",
    "        pvtol_noisy.input_labels[pvtol.ninputs:],\n",
    "    outlist = pvtol.output_labels + lqr_ctrl.output_labels,\n",
    "    outputs = pvtol.output_labels + lqr_ctrl.output_labels\n",
    ")\n",
    "print(lqr_clsys)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29f55c0a-8c17-4347-aa46-b1944e700b32",
   "metadata": {},
   "source": [
    "(The warning message can be ignored; it is generated because we implement this system as a differentially flat system and hence we require that an output function be explicitly given, rather than using `None`.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9bc481f-7b2f-4b40-89b7-1ef5a35251b7",
   "metadata": {},
   "source": [
    "We next define the characteristics of the uncertainty in the system: the disturbance and noise covariances (intensities) as well as the initial condition covariance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78853391",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Disturbance and noise intensities\n",
    "Qv = np.diag([1e-2, 1e-2])\n",
    "Qw = np.array([[1e-4, 0, 1e-5], [0, 1e-4, 1e-5], [1e-5, 1e-5, 1e-4]])\n",
    "\n",
    "# Initial state covariance\n",
    "P0 = np.eye(pvtol.nstates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c590fd88",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the time vector for the simulation\n",
    "Tf = 6\n",
    "timepts = np.linspace(0, Tf, 20)\n",
    "\n",
    "# Create representative process disturbance and sensor noise vectors\n",
    "# np.random.seed(117)           # uncomment to avoid figures changing from run to run\n",
    "V = ct.white_noise(timepts, Qv)\n",
    "W = ct.white_noise(timepts, Qw)\n",
    "plt.plot(timepts, V[0], label=\"V[0]\")\n",
    "plt.plot(timepts, W[0], label=\"W[0]\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7db5188e-03c7-439c-8cf2-47681d3feccf",
   "metadata": {},
   "source": [
    "To get a better sense of the size of the disturbances and noise, we simulate the noise-free system with the applied disturbances, and then add in the noise.  Note that in this simulation we are still assuming that the controller has access to the noise-free state (not realistic, but used here just to show that the disturbances and noise do not cause large perturbations)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c35fd695",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Desired trajectory\n",
    "xd, ud = xe, ue\n",
    "# xd = np.vstack([\n",
    "#     np.sin(2 * np.pi * timepts / timepts[-1]), \n",
    "#     np.zeros((5, timepts.size))])\n",
    "# ud = np.outer(ue, np.ones_like(timepts))\n",
    "\n",
    "# Run a simulation with full state feedback (no noise) to generate a trajectory\n",
    "uvec = [xd, ud, V, W*0]\n",
    "lqr_resp = ct.input_output_response(lqr_clsys, timepts, uvec, x0)\n",
    "U = lqr_resp.outputs[6:8]                    # controller input signals\n",
    "Y = lqr_resp.outputs[0:3] + W                # noisy output signals (noise in pvtol_noisy)\n",
    "\n",
    "# Compare to the no noise case\n",
    "uvec = [xd, ud, V*0, W*0]\n",
    "lqr0_resp = ct.input_output_response(lqr_clsys, timepts, uvec, x0)\n",
    "lqr0_fine = ct.input_output_response(lqr_clsys, timepts, uvec, x0, \n",
    "                                     t_eval=np.linspace(timepts[0], timepts[-1], 100))\n",
    "U0 = lqr0_resp.outputs[6:8]\n",
    "Y0 = lqr0_resp.outputs[0:3]\n",
    "\n",
    "# Compare the results\n",
    "# plt.plot(Y0[0], Y0[1], 'k--', linewidth=2, label=\"No disturbances\")\n",
    "plt.plot(lqr0_fine.states[0], lqr0_fine.states[1], 'r-', label=\"Actual\")\n",
    "plt.plot(Y[0], Y[1], 'b-', label=\"Noisy\")\n",
    "\n",
    "plt.xlabel('$x$ [m]')\n",
    "plt.ylabel('$y$ [m]')\n",
    "plt.axis('equal')\n",
    "plt.legend(frameon=False)\n",
    "\n",
    "plt.figure()\n",
    "plot_results(timepts, lqr_resp.states, lqr_resp.outputs[6:8])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7f1dec6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility functions for making plots\n",
    "def plot_state_comparison(\n",
    "    timepts, est_states, act_states=None, estimated_label='$\\\\hat x_{i}$', actual_label='$x_{i}$',\n",
    "    start=0):\n",
    "    for i in range(sys.nstates):\n",
    "        plt.subplot(2, 3, i+1)\n",
    "        if act_states is not None:\n",
    "            plt.plot(timepts[start:], act_states[i, start:], 'r--', \n",
    "                     label=actual_label.format(i=i))\n",
    "        plt.plot(timepts[start:], est_states[i, start:], 'b', \n",
    "                 label=estimated_label.format(i=i))\n",
    "        plt.legend()\n",
    "    plt.tight_layout()\n",
    "        \n",
    "# Define a function to plot out all of the relevant signals\n",
    "def plot_estimator_response(timepts, estimated, U, V, Y, W, start=0):\n",
    "    # Plot the input signal and disturbance\n",
    "    for i in [0, 1]:\n",
    "        # Input signal (the same across all)\n",
    "        plt.subplot(4, 3, i+1)\n",
    "        plt.plot(timepts[start:], U[i, start:], 'k')\n",
    "        plt.ylabel(f'U[{i}]')\n",
    "\n",
    "        # Plot the estimated disturbance signal\n",
    "        plt.subplot(4, 3, 4+i)\n",
    "        plt.plot(timepts[start:], estimated.inputs[i, start:], 'b-', label=\"est\")\n",
    "        plt.plot(timepts[start:], V[i, start:], 'k', label=\"actual\")\n",
    "        plt.ylabel(f'V[{i}]')\n",
    "\n",
    "    plt.subplot(4, 3, 6)\n",
    "    plt.plot(0, 0, 'b', label=\"estimated\")\n",
    "    plt.plot(0, 0, 'k', label=\"actual\")\n",
    "    plt.plot(0, 0, 'r', label=\"measured\")\n",
    "    plt.legend(frameon=False)\n",
    "    plt.grid(False)\n",
    "    plt.axis('off')\n",
    "    \n",
    "    # Plot the output (measured and estimated)   \n",
    "    for i in [0, 1, 2]:\n",
    "        plt.subplot(4, 3, 7+i)\n",
    "        plt.plot(timepts[start:], Y[i, start:], 'r', label=\"measured\")\n",
    "        plt.plot(timepts[start:], estimated.states[i, start:], 'b', label=\"measured\")\n",
    "        plt.plot(timepts[start:], Y[i, start:] - W[i, start:], 'k', label=\"actual\")\n",
    "        plt.ylabel(f'Y[{i}]')\n",
    "    \n",
    "    for i in [0, 1, 2]:\n",
    "        plt.subplot(4, 3, 10+i)\n",
    "        plt.plot(timepts[start:], estimated.outputs[i, start:], 'b', label=\"estimated\")\n",
    "        plt.plot(timepts[start:], W[i, start:], 'k', label=\"actual\")\n",
    "        plt.ylabel(f'W[{i}]')\n",
    "        plt.xlabel('Time [s]')\n",
    "\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73dd9be3",
   "metadata": {},
   "source": [
    "## State Estimation\n",
    "\n",
    "We next consider the problem of only measuring the (noisy) outputs of the system and designing a controller that uses the estimated state as the input to the LQR controller that we designed previously.\n",
    "\n",
    "We start by using a standard Kalman filter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a1f32da",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a new system with only x, y, theta as outputs\n",
    "sys = ct.nlsys(\n",
    "    pvt._noisy_update, lambda t, x, u, params: x[0:3], name=\"pvtol_noisy\",\n",
    "    states = [f'x{i}' for i in range(6)],\n",
    "    inputs = ['F1', 'F2'] + ['Dx', 'Dy'],\n",
    "    outputs = ['x', 'y', 'theta']\n",
    ")\n",
    "print(sys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a0679f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Standard Kalman filter\n",
    "linsys = sys.linearize(xe, [ue, V[:, 0] * 0])\n",
    "# print(linsys)\n",
    "B = linsys.B[:, 0:2]\n",
    "G = linsys.B[:, 2:4]\n",
    "linsys = ct.ss(\n",
    "    linsys.A, B, linsys.C, 0,\n",
    "    states=sys.state_labels, inputs=sys.input_labels[0:2], outputs=sys.output_labels)\n",
    "# print(linsys)\n",
    "\n",
    "estim = ct.create_estimator_iosystem(linsys, Qv, Qw, G=G, P0=P0)\n",
    "print(estim)\n",
    "print(f'{xe=}, {P0=}')\n",
    "\n",
    "kf_resp = ct.input_output_response(\n",
    "    estim, timepts, [Y, U], X0 = [xe, P0.reshape(-1)])\n",
    "plot_state_comparison(timepts, kf_resp.outputs, lqr_resp.states)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "654dde1b",
   "metadata": {},
   "source": [
    "### Extended Kalman filter\n",
    "\n",
    "We see that the standard Kalman filter does not do a good job in estimating the $y$ position (state $x_2$) nor the $y$ velocity (state $x_4$).\n",
    "\n",
    "A better estimate can be obtained using an extended Kalman filter, which uses the linearization of the system around the current state, rather than a fixed linearization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f83a335",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the disturbance input and measured output matrices\n",
    "F = np.array([[0, 0], [0, 0], [0, 0], [1/pvtol.params['m'], 0], [0, 1/pvtol.params['m']], [0, 0]])\n",
    "C = np.eye(3, 6)\n",
    "\n",
    "Qwinv = np.linalg.inv(Qw)\n",
    "\n",
    "# Estimator update law\n",
    "def estimator_update(t, x, u, params):\n",
    "    # Extract the states of the estimator\n",
    "    xhat = x[0:pvtol.nstates]\n",
    "    P = x[pvtol.nstates:].reshape(pvtol.nstates, pvtol.nstates)\n",
    "\n",
    "    # Extract the inputs to the estimator\n",
    "    y = u[0:3]                  # just grab the first three outputs\n",
    "    u = u[6:8]                  # get the inputs that were applied as well\n",
    "\n",
    "    # Compute the linearization at the current state\n",
    "    A = pvtol.A(xhat, u)        # A matrix depends on current state\n",
    "    # A = pvtol.A(xe, ue)       # Fixed A matrix (for testing/comparison)\n",
    "    \n",
    "    # Compute the optimal \"gain\n",
    "    L = P @ C.T @ Qwinv\n",
    "\n",
    "    # Update the state estimate\n",
    "    xhatdot = pvtol.updfcn(t, xhat, u, params) - L @ (C @ xhat - y)\n",
    "\n",
    "    # Update the covariance\n",
    "    Pdot = A @ P + P @ A.T - P @ C.T @ Qwinv @ C @ P + F @ Qv @ F.T\n",
    "\n",
    "    # Return the derivative\n",
    "    return np.hstack([xhatdot, Pdot.reshape(-1)])\n",
    "\n",
    "def estimator_output(t, x, u, params):\n",
    "    # Return the estimator states\n",
    "    return x[0:pvtol.nstates]\n",
    "\n",
    "ekf = ct.NonlinearIOSystem(\n",
    "    estimator_update, estimator_output,\n",
    "    states=pvtol.nstates + pvtol.nstates**2,\n",
    "    inputs= pvtol_noisy.output_labels \\\n",
    "        + pvtol_noisy.input_labels[0:pvtol.ninputs],\n",
    "    outputs=[f'xh{i}' for i in range(pvtol.nstates)]\n",
    ")\n",
    "print(ekf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4caf69b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ekf_resp = ct.input_output_response(\n",
    "    ekf, timepts, [lqr_resp.states, lqr_resp.outputs[6:8]],\n",
    "    X0=[xe, P0.reshape(-1)])\n",
    "plot_state_comparison(timepts, ekf_resp.outputs, lqr_resp.states)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10163c6c-5634-4dbb-ba11-e20fb1e065ed",
   "metadata": {},
   "source": [
    "## Maximum Likelihood Estimation\n",
    "\n",
    "Finally, we illustrate how to set up the problem as maximum likelihood problem, which is described in more detail in the [Optimization-Based Control](https://fbswiki.org/wiki/index.php/Supplement:_Optimization-Based_Control) (OBC) course notes, in Section 7.6.\n",
    "\n",
    "The basic idea in maximum likelihood estimation is to set up the estimation problem as an optimization problem where we define the likelihood of a given estimate (and the resulting noise and disturbances predicted by the\n",
    "model) as a cost function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1074908c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the optimal estimation problem\n",
    "traj_cost = opt.gaussian_likelihood_cost(sys, Qv, Qw)\n",
    "init_cost = lambda xhat, x: (xhat - x) @ P0 @ (xhat - x)\n",
    "oep = opt.OptimalEstimationProblem(\n",
    "        sys, timepts, traj_cost, terminal_cost=init_cost)\n",
    "\n",
    "# Compute the estimate from the noisy signals\n",
    "est = oep.compute_estimate(Y, U, X0=lqr_resp.states[:, 0])\n",
    "plot_state_comparison(timepts, est.states, lqr_resp.states)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c6981b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the response of the estimator\n",
    "plot_estimator_response(timepts, est, U, V, Y, W)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25b8aa85",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Noise free and disturbance free => estimation should be near perfect\n",
    "noisefree_cost = opt.gaussian_likelihood_cost(sys, Qv, Qw*1e-6)\n",
    "oep0 = opt.OptimalEstimationProblem(\n",
    "        sys, timepts, noisefree_cost, terminal_cost=init_cost)\n",
    "est0 = oep0.compute_estimate(Y0, U0, X0=lqr0_resp.states[:, 0],\n",
    "                            initial_guess=(lqr0_resp.states, V * 0))\n",
    "plot_state_comparison(\n",
    "    timepts, est0.states, lqr0_resp.states, estimated_label='$\\\\bar x_{i}$')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a76821f",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_estimator_response(timepts, est0, U0, V*0, Y0, W*0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b9031cf",
   "metadata": {},
   "source": [
    "### Bounded disturbances\n",
    "\n",
    "Another situation that the maximum likelihood framework can handle is when input distributions that are bounded.  We implement that here by carrying out the optimal estimation problem with constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93482470",
   "metadata": {},
   "outputs": [],
   "source": [
    "V_clipped = np.clip(V, -0.05, 0.05)    \n",
    "\n",
    "plt.plot(timepts, V[0], label=\"V[0]\")\n",
    "plt.plot(timepts, V_clipped[0], label=\"V[0] clipped\")\n",
    "plt.plot(timepts, W[0], label=\"W[0]\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56e186f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "uvec = [xe, ue, V_clipped, W]\n",
    "clipped_resp = ct.input_output_response(lqr_clsys, timepts, uvec, x0)\n",
    "U_clipped = clipped_resp.outputs[6:8]        # controller input signals\n",
    "Y_clipped = clipped_resp.outputs[0:3] + W    # noisy output signals\n",
    "\n",
    "traj_constraint = opt.disturbance_range_constraint(\n",
    "    sys, [-0.05, -0.05], [0.05, 0.05])\n",
    "oep_clipped = opt.OptimalEstimationProblem(\n",
    "        sys, timepts, traj_cost, terminal_cost=init_cost,\n",
    "        trajectory_constraints=traj_constraint)\n",
    "\n",
    "est_clipped = oep_clipped.compute_estimate(\n",
    "    Y_clipped, U_clipped, X0=lqr0_resp.states[:, 0])\n",
    "plot_state_comparison(timepts, est_clipped.states, lqr_resp.states)\n",
    "plt.suptitle(\"MHE with constraints\")\n",
    "plt.tight_layout()\n",
    "\n",
    "plt.figure()\n",
    "ekf_unclipped = ct.input_output_response(\n",
    "    ekf, timepts, [clipped_resp.states, clipped_resp.outputs[6:8]],\n",
    "    X0=[xe, P0.reshape(-1)])\n",
    "\n",
    "plot_state_comparison(timepts, ekf_unclipped.outputs, lqr_resp.states)\n",
    "plt.suptitle(\"EKF w/out constraints\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "108c341a",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_estimator_response(timepts, est_clipped, U, V_clipped, Y, W)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "430117ce",
   "metadata": {},
   "source": [
    "## Moving Horizon Estimation (MHE)\n",
    "\n",
    "Finally, we can now move to the implementation of a moving horizon estimator, using our fixed horizon, maximum likelihood, optimal estimator.  The details of this implementation are described in more detail in the [Optimization-Based Control](https://fbswiki.org/wiki/index.php/Supplement:_Optimization-Based_Control) (OBC) course notes, in Section 7.6."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "121d67ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use a shorter horizon\n",
    "mhe_timepts = timepts[0:5]\n",
    "oep = opt.OptimalEstimationProblem(\n",
    "        sys, mhe_timepts, traj_cost, terminal_cost=init_cost)\n",
    "\n",
    "try:\n",
    "    mhe = oep.create_mhe_iosystem(2)\n",
    "    \n",
    "    est_mhe = ct.input_output_response(\n",
    "        mhe, timepts, [Y, U], X0=resp.states[:, 0], \n",
    "        params={'verbose': True}\n",
    "    )\n",
    "    plot_state_comparison(timepts, est_mhe.states, lqr_resp.states)\n",
    "except:\n",
    "    print(\"MHE for continuous-time systems not implemented\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1914ad96",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create discrete-time version of PVTOL\n",
    "Ts = 0.1\n",
    "print(f\"Sample time: {Ts=}\")\n",
    "dsys = ct.nlsys(\n",
    "    lambda t, x, u, params: x + Ts * sys.updfcn(t, x, u, params),\n",
    "    sys.outfcn, dt=Ts, states=sys.state_labels,\n",
    "    inputs=sys.input_labels, outputs=sys.output_labels,\n",
    ")\n",
    "print(dsys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11162130",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a new list of time points\n",
    "timepts = np.arange(0, Tf, Ts)\n",
    "\n",
    "# Create representative process disturbance and sensor noise vectors\n",
    "# np.random.seed(117)           # avoid figures changing from run to run\n",
    "V = ct.white_noise(timepts, Qv)\n",
    "# V = np.clip(V0, -0.1, 0.1)    # Hold for later\n",
    "W = ct.white_noise(timepts, Qw, dt=Ts)\n",
    "# plt.plot(timepts, V0[0], 'b--', label=\"V[0]\")\n",
    "plt.plot(timepts, V[0], label=\"V[0]\")\n",
    "plt.plot(timepts, W[0], label=\"W[0]\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8a6a693",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate a new trajectory over the longer time vector\n",
    "uvec = [xd, ud, V, W*0]\n",
    "lqr_resp = ct.input_output_response(lqr_clsys, timepts, uvec, x0)\n",
    "U = lqr_resp.outputs[6:8]                    # controller input signals\n",
    "Y = lqr_resp.outputs[0:3] + W                # noisy output signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d683767f",
   "metadata": {},
   "outputs": [],
   "source": [
    "mhe_timepts = timepts[0:10]\n",
    "oep = opt.OptimalEstimationProblem(\n",
    "        dsys, mhe_timepts, traj_cost, terminal_cost=init_cost,\n",
    "        disturbance_indices=[2, 3])\n",
    "mhe = oep.create_mhe_iosystem()\n",
    "    \n",
    "mhe_resp = ct.input_output_response(\n",
    "    mhe, timepts, [Y, U], X0=x0, \n",
    "    params={'verbose': True}\n",
    ")\n",
    "plot_state_comparison(timepts, mhe_resp.states, lqr_resp.states)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad6aac39-5b55-4ffd-ab21-44385dc11ff5",
   "metadata": {},
   "source": [
    "Although this estimator eventually converges to the underlying tate of the system, the initial transient response is quite poor.\n",
    "\n",
    "One possible explanation is that we are not starting the system at the origin, even though we are penalizing the initial state if it is away from the origin.\n",
    "\n",
    "To see if this matters, we shift the problem to one in which the system's initial condition is at the origin:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfc68072",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Resimulate starting at the origin and moving to the \"initial\" condition\n",
    "uvec = [x0, ue, V, W*0]\n",
    "lqr_resp = ct.input_output_response(lqr_clsys, timepts, uvec, xe)\n",
    "U = lqr_resp.outputs[6:8]                    # controller input signals\n",
    "Y = lqr_resp.outputs[0:3] + W                # noisy output signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49213d04",
   "metadata": {},
   "outputs": [],
   "source": [
    "mhe_timepts = timepts[0:8]\n",
    "oep = opt.OptimalEstimationProblem(\n",
    "        dsys, mhe_timepts, traj_cost, terminal_cost=init_cost,\n",
    "        disturbance_indices=[2, 3])\n",
    "mhe = oep.create_mhe_iosystem()\n",
    "    \n",
    "mhe_resp = ct.input_output_response(\n",
    "    mhe, timepts, [Y, U],\n",
    "    params={'verbose': True}\n",
    ")\n",
    "plot_state_comparison(timepts, mhe_resp.outputs, lqr_resp.states)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "650a559a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
